import warnings
warnings.filterwarnings('ignore')
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import tensorflow as tf
tf.compat.v1.logging.set_verbosity(40)

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Activation, Flatten, Dropout

from tensorflow.keras.layers import Conv2D, MaxPooling2D
from tensorflow.keras import losses, optimizers, metrics

import os
import cv2
import numpy as np

# data sets
IMG_HEIGHT = 60
IMG_WIDTH = 160
char_num = 4
characters = range(10)
labellen = char_num * len(characters)
print("labellen", labellen)

batch_size = 100
nb_epoch = 10

"""
验证码标签为40维的向量 label: 1327
    [[0,1,0,0,0,0,0,0,0,0],
    [0,0,0,1,0,0,0,0,0,0],
    [0,0,1,0,0,0,0,0,0,0],
    [0,0,0,0,0,0,0,1,0,0]]
"""
def label2vec(label):
    label_vec = np.zeros(labellen)
    for i, num in enumerate(label):
        idx = i * len(characters) + int(num)
        label_vec[idx] = 1
    return label_vec

def readData(file_path):
    x_images = []
    y_labels = []
    for item in os.listdir(file_path):
        if '.nomedia' == item:
            continue
        item_path = file_path + '/' + item
        image = cv2.imread(item_path) / 255
        x_images.append(image)
        label = os.path.splitext(item)[0]
        y_labels.append(label2vec(label))
    return np.array(x_images), np.array(y_labels)

test_dir = r"../../../../large_data/DL1/_many_files/vcode_data/train"
x_images, y_labels = readData(test_dir)

from sklearn.model_selection import train_test_split

x_train, x_test, y_train, y_test = train_test_split(x_images, y_labels, train_size=0.9)

model = Sequential([
    Conv2D(6, (5, 5), activation='relu'),
    MaxPooling2D((2, 2)),
    Conv2D(16, (5, 5), activation='relu'),
    MaxPooling2D((2, 2)),

    Flatten(),

    Dense(120, activation='relu'),
    Dense(84, activation='relu'),
    Dense(labellen, activation='sigmoid', name="logits")
])

model.compile(loss=losses.binary_crossentropy,
              optimizer=optimizers.Adam(0.0001),
              metrics=['accuracy'])

model.fit(x_train, y_train, batch_size=batch_size, epochs=nb_epoch)

score = model.evaluate(x_test, y_test)
print('test loss:', score[0])
print('test acc:', score[1])
